import numpy as np
import matplotlib.pyplot as plt


def PE(vecs):
    length, dim = vecs.shape
    pos = np.tile(np.arange(length).reshape(-1,1), (1, dim))
    print("pos.shape:", pos.shape)
    d_indice = np.arange(dim)
    d_mask_even = np.array([1 if i%2 == 0 else 0 for i in d_indice])
    d_mask_odd = np.array([1 if i%2 == 1 else 0 for i in d_indice])
    i_list = np.power(10000, d_indice/dim)
    pe = np.sin(pos / i_list) * d_mask_even + np.cos(pos / i_list) * d_mask_odd
    print("pe.shape:", pe.shape)
    return pe


def plot(pos_encoding):
    new_axis, tokens, dimensions = pos_encoding.shape
    plt.figure(figsize=(12, 8))
    plt.pcolormesh(pos_encoding[0], cmap='viridis')
    plt.xlabel('Embedding Dimensions')
    plt.xlim((0, dimensions))
    plt.ylim((tokens, 0))
    plt.ylabel('Token Position')
    plt.colorbar()
    plt.show()


def main():
    vecs = np.random.randn(10,30)
    # print(vecs.shape)
    r = PE(vecs)
    print(r.shape)
    plot(r[np.newaxis, ...])
    # print(r[1])
    # plt.plot(np.arange(vecs.shape[1]), r[1])
    # plt.plot(np.arange(vecs.shape[1]), r[9])
    # plt.show()


def sample_from_zhihu():
    import numpy as np
    import matplotlib.pyplot as plt

    # Code from https://www.tensorflow.org/tutorials/text/transformer
    def get_angles(pos, i, d_model):
        angle_rates = 1 / np.power(10000, (2 * (i // 2)) / np.float32(d_model))
        return pos * angle_rates

    def positional_encoding(position, d_model):
        angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                                np.arange(d_model)[np.newaxis, :],
                                d_model)

        # apply sin to even indices in the array; 2i
        angle_rads[:, 0::2] = np.sin(angle_rads[:, 0::2])

        # apply cos to odd indices in the array; 2i+1
        angle_rads[:, 1::2] = np.cos(angle_rads[:, 1::2])

        pos_encoding = angle_rads[np.newaxis, ...]

        return pos_encoding

    tokens = 10
    dimensions = 64

    pos_encoding = positional_encoding(tokens, dimensions)
    print(pos_encoding.shape)

    plt.figure(figsize=(12, 8))
    plt.pcolormesh(pos_encoding[0], cmap='viridis')
    plt.xlabel('Embedding Dimensions')
    plt.xlim((0, dimensions))
    plt.ylim((tokens, 0))
    plt.ylabel('Token Position')
    plt.colorbar()
    plt.show()

if __name__ == "__main__":
    main()
    # sample_from_zhihu()
